# @Time : 2025/4/19 1:21
# @Author : ZhangQi
# adb_agent.py
import inspect
import json
import subprocess
import time
from typing import Optional, Dict, Any, List
from openai import OpenAI

from config import ALIYUN_CONFIG


class ADBAgent:
    def __init__(self, device_id: Optional[str] = None):
        self.device_id = device_id
        self.client = OpenAI(
            api_key=ALIYUN_CONFIG.get("aliyun_bailian_api_key"),
            base_url=ALIYUN_CONFIG.get("dashscope_compatible_base_url")
        )
        self.available_functions = self._get_available_functions()
        print("已注册功能列表:", self._list_registered_functions())

    def _get_available_functions(self) -> Dict[str, Any]:
        return {
            name: func
            for name, func in inspect.getmembers(self, inspect.ismethod)
            if hasattr(func, '_function_schema')
        }

    def _list_registered_functions(self) -> List[str]:
        return [func._function_schema["function"]["name"] for func in self.available_functions.values()]

    def execute_adb(self, command: str) -> str:
        """ADB命令执行（编码问题修复版）"""
        try:
            # 构建完整命令
            full_cmd = ["adb"]
            if self.device_id:
                full_cmd.extend(["-s", self.device_id])
            full_cmd += command.split()

            # 执行命令并处理编码
            result = subprocess.run(
                full_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                check=True,
                timeout=30,
                encoding='utf-8',  # 显式指定编码
                errors='replace'  # 替换无法解码的字符
            )
            return result.stdout.strip()

        except subprocess.CalledProcessError as e:
            error_msg = e.output.strip() if e.output else "Unknown error"
            raise RuntimeError(f"ADB命令失败: {error_msg}") from e

    def _take_emergency_screenshot(self):
        """紧急情况截屏"""
        timestamp = int(time.time())
        self.execute_adb(f"shell screencap -p /sdcard/emergency_{timestamp}.png")
        self.execute_adb(f"pull /sdcard/emergency_{timestamp}.png .")

    def process_command(self, natural_language: str) -> Dict:
        """处理自然语言指令（阿里云适配版）"""
        tools = [func._function_schema for func in self.available_functions.values()]

        try:
            response = self.client.chat.completions.create(
                model="qwen-plus",
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个ADB自动化助手，必须使用JSON格式调用工具函数。"
                                   "严格按照以下步骤执行："
                                   "1. 分析用户需求 2. 选择合适工具 3. 生成正确参数"
                    },
                    {"role": "user", "content": natural_language}
                ],
                tools=tools,
                extra_body={
                    "plugin_type": "tools_call",
                    "plugins": {"tools_call": {"enabled": True}}
                }
            )
        except Exception as e:
            return {"status": "error", "message": str(e)}

        return self._parse_response(response)

    def _parse_response(self, response) -> Dict:
        """解析阿里云响应"""
        message = response.choices[0].message
        if not message.tool_calls:
            return {"status": "no_action", "content": message.content}

        results = []
        for tool in message.tool_calls:
            func_name = tool.function.name
            try:
                kwargs = json.loads(tool.function.arguments)
                func = self._get_function_by_name(func_name)
                result = func(**kwargs)
                results.append({
                    "function": func_name,
                    "arguments": kwargs,
                    "result": result,
                    "status": "success"
                })
            except Exception as e:
                results.append({
                    "function": func_name,
                    "error": str(e),
                    "status": "failed"
                })
        return {"results": results}

    def _get_function_by_name(self, name: str):
        for func in self.available_functions.values():
            if func._function_schema["function"]["name"] == name:
                return func
        raise ValueError(f"未找到函数: {name}")
